# mindmap.py
import os,logger
import json
import shutil
import requests
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.Qt import QTextOption

class MindMapModule(QWidget):
    content_updated = pyqtSignal(dict)
    
    def __init__(self, main_window):
        super().__init__()
        self.main_window = main_window
        self.essays = []
        self.current_essay = None
        self.data_dir = os.path.abspath("data/essays")
        self.api_key = ""
        os.makedirs(self.data_dir, exist_ok=True)
        
        self.init_ui()
        self.load_essays()

    def init_ui(self):
        main_layout = QHBoxLayout(self)
        main_layout.setContentsMargins(15, 15, 15, 15)
        main_layout.setSpacing(15)
        
        # 左侧列表区域
        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)
        
        # 列表控件
        self.list_widget = QListWidget()
        self.list_widget.setMinimumWidth(200)
        self.list_widget.setStyleSheet("""
            QListWidget {
                background: white;
                border: 1px solid #ddd;
                border-radius: 4px;
            }
            QListWidget::item {
                padding: 6px;
            }
            QListWidget::item:selected {
                background: #2196F3;
                color: white;
            }
        """)
        self.list_widget.itemClicked.connect(self.load_essay)
        
        # 控制按钮组
        btn_group = QWidget()
        btn_layout = QHBoxLayout(btn_group)
        btn_layout.setContentsMargins(0, 0, 0, 0)
        btn_layout.setSpacing(5)
        
        self.new_btn = self.create_tool_button("新建", "icons/add.svg", "#4CAF50")
        self.del_btn = self.create_tool_button("删除", "icons/delete.svg", "#f44336")
        self.save_btn = self.create_tool_button("保存", "icons/save.svg", "#2196F3")
        
        btn_layout.addWidget(self.new_btn)
        btn_layout.addWidget(self.del_btn)
        btn_layout.addWidget(self.save_btn)
        
        left_layout.addWidget(btn_group)
        left_layout.addWidget(self.list_widget)
        
        # 右侧编辑区域
        right_panel = QWidget()
        right_layout = QVBoxLayout(right_panel)
        right_layout.setSpacing(10)
        
        # 选项卡
        self.tabs = QTabWidget()
        self.tabs.setStyleSheet("""
            QTabBar::tab {
                min-width: 100px;
                padding: 8px 16px;
            }
        """)
        
        # 原文编辑区
        self.original_edit = QTextEdit()
        self.original_edit.setStyleSheet("""
            QTextEdit {
                background: white;
                color: #333;
                border: 1px solid #ddd;
                border-radius: 4px;
                padding: 12px;
                font-size: 14px;
            }
        """)
        self.original_edit.setWordWrapMode(QTextOption.WrapAtWordBoundaryOrAnywhere)
        
        # 思维导图区
        self.argument_tree = self.create_tree_widget()
        self.logic_tree = self.create_tree_widget()
        
        self.tabs.addTab(self.original_edit, "原文编辑")
        self.tabs.addTab(self.argument_tree, "论点结构")
        self.tabs.addTab(self.logic_tree, "逻辑分析")
        
        # 生成按钮
        self.gen_btn = QPushButton(" 生成思维导图")
        self.gen_btn.setIcon(QIcon("icons/mindmap.svg"))
        self.gen_btn.setStyleSheet(self.get_button_style("#9C27B0"))
        self.gen_btn.setCursor(Qt.PointingHandCursor)
        
        right_layout.addWidget(self.tabs, 1)
        right_layout.addWidget(self.gen_btn)
        
        main_layout.addWidget(left_panel, 1)
        main_layout.addWidget(right_panel, 3)
        
        # 信号连接
        self.new_btn.clicked.connect(self.new_essay)
        self.del_btn.clicked.connect(self.delete_essay)
        self.save_btn.clicked.connect(self.save_current_essay)
        self.gen_btn.clicked.connect(self.generate_mindmap)
        
        self.apply_styles()

    def create_tool_button(self, text, icon_path, color):
        btn = QPushButton(text)
        btn.setIcon(QIcon(icon_path))
        btn.setStyleSheet(f"""
            QPushButton {{
                padding: 8px 12px;
                border-radius: 4px;
                background: {color};
                color: white;
                border: none;
                min-width: 80px;
            }}
            QPushButton:hover {{
                background: {color}DD;
            }}
            QPushButton:pressed {{
                background: {color}BB;
            }}
        """)
        btn.setCursor(Qt.PointingHandCursor)
        return btn

    def apply_styles(self):
        self.setStyleSheet("""
            QTreeWidget, QListWidget {
                border: 1px solid #e0e0e0;
                border-radius: 8px;
                padding: 8px;
            }
            QPushButton {
                padding: 8px;
                border-radius: 4px;
                min-width: 80px;
            }
        """)
    
    def get_button_style(self, color):
        return f"""
            QPushButton {{
                background: {color};
                color: white;
                padding: 10px 20px;
                border-radius: 4px;
                border: none;
                min-width: 120px;
            }}
            QPushButton:hover {{
                background: {color}dd;
            }}
        """

    def create_tree_widget(self):
        tree = QTreeWidget()
        tree.setHeaderHidden(True)
        tree.setStyleSheet("""
            QTreeWidget {
                background: white;
                border: 1px solid #ddd;
                border-radius: 4px;
            }
            QTreeWidget::item {
                padding: 6px;
            }
        """)
        return tree

    def load_essays(self):
        self.list_widget.clear()
        has_items = False
        try:
            for item in os.listdir(self.data_dir):
                path = os.path.join(self.data_dir, item)
                if os.path.isdir(path):
                    self.list_widget.addItem(item)
                    has_items = True
            if not has_items:
                QMessageBox.information(self, "提示", "作文目录为空，请新建作文")
        except FileNotFoundError:
            QMessageBox.critical(self, "错误", f"数据目录不存在：{self.data_dir}")
        except PermissionError:
            QMessageBox.critical(self, "权限错误", f"无法访问目录：{self.data_dir}")
    
    def show_context_menu(self, pos):
        item = self.list_widget.itemAt(pos)
        if item:
            self.context_menu.exec_(self.list_widget.mapToGlobal(pos))

    def update_api_key(self, key):
        self.api_key = key


    def delete_essay(self):
        current_item = self.list_widget.currentItem()
        if current_item:
            essay_name = current_item.text()
            essay_path = os.path.join(self.data_dir, essay_name)
            if os.path.exists(essay_path):
                shutil.rmtree(essay_path)
                self.list_widget.takeItem(self.list_widget.row(current_item))
                QMessageBox.information(self, "成功", "作文已删除！")

    def sanitize_filename(self, name):
        invalid_chars = '<>:"/\\|?*'
        replace_char = '_'
        cleaned = ''.join([c if c not in invalid_chars else replace_char for c in name])
        return cleaned.strip()[:50]

    def new_essay(self):
        title, ok = QInputDialog.getText(self, "新建作文", "请输入作文标题：")
        if ok and title:
            try:
                clean_title = self.sanitize_filename(title)
                essay_path = os.path.join(self.data_dir, clean_title)
                
                if os.path.exists(essay_path):
                    raise FileExistsError(f"已存在同名作文：{clean_title}")
                
                # 创建目录和空文件
                os.makedirs(essay_path, exist_ok=True)
                with open(os.path.join(essay_path, "original.txt"), 'w', encoding='utf-8') as f:
                    f.write("")
                
                self.list_widget.addItem(clean_title)
                self.current_essay = clean_title
                self.original_edit.clear()
                QMessageBox.information(self, "成功", f"作文 [{clean_title}] 创建成功！")
            except Exception as e:
                QMessageBox.critical(self, "错误", f"创建失败：\n{str(e)}")

    def load_essay(self, item):
        if not item:
            return
            
        essay_name = item.text()
        essay_path = os.path.join(self.data_dir, essay_name)
        original_path = os.path.join(essay_path, "original.txt")
        
        try:
            # 检查文件是否存在
            if not os.path.exists(original_path):
                raise FileNotFoundError(f"找不到文件：{original_path}")
            
            # 读取内容
            with open(original_path, "r", encoding="utf-8") as f:
                self.original_edit.setPlainText(f.read())
            
            self.current_essay = essay_name
            
            # 加载分析数据
            analysis_types = {
                "argument": self.argument_tree,
                "logic": self.logic_tree
            }
            
            for atype, tree in analysis_types.items():
                analysis_path = os.path.join(essay_path, f"{atype}_analysis.json")
                if os.path.exists(analysis_path):
                    with open(analysis_path, "r", encoding="utf-8") as f:
                        raw_data = f.read()
                        # 新增：解析Unicode转义字符
                        data = json.loads(raw_data.encode('utf-8').decode('unicode_escape'))
                        self.update_tree(tree, data)
                        
                else:
                    tree.clear()
                    
        except Exception as e:
            QMessageBox.critical(self, "加载错误",
                f"无法加载作文：{essay_name}\n"
                f"错误原因：{str(e)}\n"
                "建议操作：\n1. 尝试重新生成\n2. 删除并新建作文")

    def handle_api_error(self, error):
        QMessageBox.critical(self, "错误", f"API调用失败：\n{error}")

    def save_current_essay(self):

        if not self.current_essay:
            QMessageBox.warning(self, "警告", "请先选择或创建作文！")
            return
        
        essay_path = os.path.join(self.data_dir, self.current_essay)
        original_path = os.path.join(essay_path, "original.txt")
        # 增加列表项存在性检查
        items = self.list_widget.findItems(self.current_essay, Qt.MatchExactly)
        if not items:
            QMessageBox.critical(self, "错误", "当前作文不存在于列表中！")
            return
        try:
            essay_path = os.path.join(self.data_dir, self.current_essay)
            os.makedirs(essay_path, exist_ok=True)  # 确保目录存在
            
            # 增强文件路径验证
            original_path = os.path.join(essay_path, "original.txt")
            if not os.path.exists(os.path.dirname(original_path)):
                raise FileNotFoundError("存储路径不存在")
    
            # 保持原子写入逻辑
            temp_path = os.path.join(essay_path, "temp_original.txt")
            with open(temp_path, "w", encoding="utf-8") as f:
                f.write(self.original_edit.toPlainText())
    
            # 替换文件前验证可写性
            if os.path.exists(original_path):
                if not os.access(original_path, os.W_OK):
                    raise PermissionError("文件不可写")
                os.remove(original_path)
                
            os.rename(temp_path, original_path)
            
            # 保存分析数据（添加空值保护）
            if self.argument_tree.topLevelItemCount() > 0 or self.logic_tree.topLevelItemCount() > 0:
                self.save_analysis()
            else:
                QMessageBox.information(self, "提示", "已保存原文，未检测到分析数据")
                
            QMessageBox.information(self, "保存成功", "所有数据已安全存储！")
            
        except Exception as e:
            QMessageBox.critical(self, "保存失败", f"错误详情：\n{str(e)}")

    def generate_mindmap(self):
        if not self.current_essay:
            QMessageBox.warning(self, "警告", "请先选择或创建作文！")
            return
            
        if not self.api_key:
            QMessageBox.critical(self, "错误", "请先在设置中配置API密钥！")
            return
            
        content = self.original_edit.toPlainText()
        if len(content.strip()) < 50:
            QMessageBox.warning(self, "内容过短", "作文内容至少需要50个字符！")
            return
            
        try:
            self.worker = APIWorker(self.api_key, content)
            self.worker.finished.connect(self.handle_api_response)
            self.worker.error.connect(self.handle_api_error)
            self.worker.start()
            QMessageBox.information(self, "开始处理", "正在生成思维导图，请稍候...")
        except Exception as e:
            QMessageBox.critical(self, "启动失败", f"无法开始处理：{str(e)}")

    def handle_api_response(self, result):
        try:
            argument_data, logic_data = result
            
            # 验证数据结构
            if not all(key in argument_data.get("structure", {}) for key in ["name", "children"]):
                raise ValueError("论点结构数据格式错误")
            if not all(key in logic_data.get("structure", {}) for key in ["name", "children"]):
                raise ValueError("逻辑分析数据格式错误")
                
            # 更新视图
            self.update_tree(self.argument_tree, argument_data)
            self.update_tree(self.logic_tree, logic_data)
            
            # 自动保存
            self.save_analysis()
            QMessageBox.information(self, "生成成功", "思维导图已生成并保存！")
        except Exception as e:
            QMessageBox.critical(self, "数据处理错误", f"{str(e)}\n建议检查API密钥或联系开发者")

    # def update_tree(self, tree, data):
    #     tree.clear()
    #     try:
    #         if isinstance(data, dict) and "structure" in data:
    #             root = QTreeWidgetItem(tree, [data["structure"].get("name", "根节点")])
    #             self.build_tree(root, data["structure"].get("children", []))
    #             tree.expandAll()
    #     except Exception as e:
    #         print(f"树形结构加载失败: {str(e)}")

    def update_tree(self, tree, data):
        tree.clear()
        try:
            # 1. 检查是否为有效字典
            if not isinstance(data, dict):
                raise ValueError("数据必须为字典类型")

            # 2. 提取结构数据（带默认值）
            structure = data.get("structure", {})
            if not isinstance(structure, dict):
                structure = {}  # 重置为默认值

            # 3. 提取根节点名称（强制字符串转换）
            root_name = str(structure.get("name", "未命名根节点"))

            # 4. 创建根节点
            root = QTreeWidgetItem(tree, [root_name])

            # 5. 提取子节点（强制列表类型）
            children = structure.get("children", [])
            if not isinstance(children, list):
                children = []  # 非法类型时重置为空列表

            # 6. 递归构建子树
            self.build_tree(root, children)

            # 7. 自动展开
            tree.expandAll()

        except Exception as e:
            logger.error(f"树形结构加载失败: {str(e)}")
            QMessageBox.critical(
                self, "数据错误",
                f"无法加载思维导图：\n{str(e)}\n"
                "请检查JSON文件格式是否符合要求。"
            )

    def build_tree(self, parent, children):
        for child in children:
            try:
                # 1. 跳过非字典类型的子节点
                if not isinstance(child, dict):
                    logger.warning(f"跳过无效子节点（类型错误）: {child}")
                    continue

                # 2. 提取节点名称（强制字符串转换）
                node_name = str(child.get("name", "未命名节点"))

                # 3. 创建节点
                node = QTreeWidgetItem(parent, [node_name])

                # 4. 提取子节点列表（强制列表类型）
                sub_children = child.get("children", [])
                if not isinstance(sub_children, list):
                    sub_children = []
                    logger.warning(f"节点 '{node_name}' 的子节点类型错误，已重置为空列表")

                # 5. 递归构建
                self.build_tree(node, sub_children)

            except Exception as e:
                logger.error(f"构建子树失败: {str(e)}")
                continue





    def save_analysis(self):
        essay_path = os.path.join(self.data_dir, self.current_essay)
        for analysis_type, tree in [("argument", self.argument_tree), ("logic", self.logic_tree)]:
            data = self.tree_to_dict(tree)
            with open(os.path.join(essay_path, f"{analysis_type}_analysis.json"), "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2)

    def tree_to_dict(self, tree):
        def parse_item(item):
            if not item:
                return {}
            return {
                "name": item.text(0),
                "children": [parse_item(child) for i in range(item.childCount()) 
                            for child in [item.child(i)] if child]
            }

        # 添加根节点存在性检查
        root = tree.invisibleRootItem()
        if root.childCount() == 0:
            return {"structure": {"name": "未命名", "children": []}}

        return {"structure": parse_item(root.child(0))}

class APIWorker(QThread):
    finished = pyqtSignal(tuple)
    error = pyqtSignal(str)

    def __init__(self, api_key, content):
        super().__init__()
        self.api_key = api_key
        self.content = content
        self.retry_count = 3

    def run(self):
        for attempt in range(self.retry_count):
            try:
                argument = self.call_api("argument")
                logic = self.call_api("logic")
                self.finished.emit((argument, logic))
                return
            except Exception as e:
                if attempt == self.retry_count - 1:
                    self.error.emit(f"API调用失败（尝试{self.retry_count}次）：{str(e)}")
                else:
                    self.msleep(1000)

    def call_api(self, analysis_type):
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        prompt = f"""请严格按以下层级结构分析作文：
要求：
0. 论据必须详细
1. 必须使用严格的JSON格式
2. 每个节点必须是对象且包含name字段
3. 子节点必须放在children数组中
4. 分析类型：{'论点结构（中心论点-分论点-论据）' if analysis_type == 'argument' else '行文逻辑与辩证分析'}
返回格式示例：
{{
    "structure": {{
        "name": "根节点名称",
        "children": [
            {{"name": "子节点1", "children": [...]}},
            {{"name": "子节点2"}}
        ]
    }}
}}
作文内容：{self.content}"""
        
        data = {
            "model": "deepseek-reasoner",
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.3,
            "max_tokens": 2000
        }
        
        try:
            response = requests.post(
                "https://api.deepseek.com/chat/completions",
                headers=headers,
                json=data,
                # timeout=(10, 30)
            )
            response.raise_for_status()
            
            content = response.json()["choices"][0]["message"]["content"]
            json_str = content.replace('```json', '').replace('```', '').strip()
            return json.loads(json_str)
            
        except requests.exceptions.RequestException as e:
            raise Exception(f"网络错误：{str(e)}")
        except json.JSONDecodeError:
            raise Exception("API返回了无效的JSON格式")
        except KeyError:
            raise Exception("API响应结构不符合预期")